from jinja2 import Environment, FileSystemLoader
import os
from dsl_parser import (
    parse_dsl,
    BasicType,
    CustomizedType,
    VectorType,
    MapType,
    MemberDecl,
    ClassType,
    StructType,
    EnumItem,
    EnumType,
    ArrayType,
)
import argparse
import os
import pprint
from typing import Union


BASIC_TYPE_TO_CPP_TYPE = {
    "string": "std::string",
    "float32_t": "float",
    "float64_t": "double",
}

SMALL_ALGORITHMIC_TYPE = set(
    [
        "int8_t",
        "uint8_t",
        "int16_t",
        "uint16_t",
        "int32_t",
        "uint32_t",
        "int64_t",
        "uint64_t",
        "float32_t",
        "float64_t",
        "bool",
    ]
)


def get_message_type(type: Union[StructType, ClassType, EnumType]) -> str:
    if isinstance(type, StructType):
        return "struct"
    if isinstance(type, ClassType):
        return "class"
    if isinstance(type, EnumType):
        return "enum"
    raise ValueError(f"Unknown message type {type}")


def get_cpp_type_name(
    T: Union[BasicType, CustomizedType, "VectorType", "MapType", "ArrayType"]
) -> str:
    if isinstance(T, BasicType):
        if T.type in BASIC_TYPE_TO_CPP_TYPE:
            return BASIC_TYPE_TO_CPP_TYPE[T.type]
        return T.type
    elif isinstance(T, CustomizedType):
        return T.type.replace(".", "::")
    elif isinstance(T, VectorType):
        return f"std::vector<{get_cpp_type_name(T.type)}>"
    elif isinstance(T, MapType):
        return f"std::map<{get_cpp_type_name(T.key_type)}, {get_cpp_type_name(T.value_type)}>"
    elif isinstance(T, ArrayType):
        return f"std::array<{get_cpp_type_name(T.type)}, {T.num}>"
    raise ValueError(f"Unknow type {T}")


def get_return_type(T: Union[BasicType, CustomizedType, "VectorType", "MapType"]):
    if isinstance(T, BasicType):
        if T.type in SMALL_ALGORITHMIC_TYPE:
            return get_cpp_type_name(T)
    return f"const {get_cpp_type_name(T)} &"


def should_use_move(T):
    if isinstance(T, BasicType):
        if T.type in SMALL_ALGORITHMIC_TYPE:
            return False
    return True


def get_move_statement(field) -> str:
    T = field.type
    name = field.name
    if should_use_move(T):
        return f"std::move({name})"
    return name


def get_initializer_list(class_type: ClassType) -> str:
    init_list = []
    for field in class_type.fields:
        T = field.type
        name = field.name
        if should_use_move(T):
            init_list.append(f"m_{name}(std::move({name}))")
        else:
            init_list.append(f"m_{name}({name})")
    return ",".join(init_list)


def get_cntr_signature(class_type: ClassType):
    def get_constructor_list(class_type) -> str:
        arguments = []
        for field in class_type.fields:
            T = field.type
            name = field.name
            arguments.append(f"{get_cpp_type_name(T)} {name}")
        return ", ".join(arguments)

    result = ""
    if len(class_type.fields) == 1:
        result = "explicit "

    result += f"{class_type.name}({get_constructor_list(class_type)})"
    return result


def generate_cpp_code(schema_file, output_dir):
    with open(schema_file) as f:
        schema_content = f.read()

    dsl_definitions = {"messages": parse_dsl(schema_content)}
    env = Environment(
        loader=FileSystemLoader("dsl/"), trim_blocks=True, lstrip_blocks=True
    )
    env.filters["get_message_type"] = get_message_type
    env.filters["get_cpp_type_name"] = get_cpp_type_name
    env.filters["get_return_type"] = get_return_type
    env.filters["get_move_statement"] = get_move_statement
    env.filters["get_initializer_list"] = get_initializer_list
    env.filters["get_cntr_signature"] = get_cntr_signature

    template = env.get_template("template.h.j2")
    rendered = template.render(messages=dsl_definitions["messages"])

    os.makedirs(output_dir, exist_ok=True)
    schema_basename = os.path.basename(schema_file).split(".")[0]

    output_path = os.path.join(output_dir, schema_basename + ".hpp")

    with open(output_path, "w") as f:
        f.write(rendered)

    print(f"Generated code saved to {output_path}")


if __name__ == "__main__":
    print("----------------------------------------")
    parser = argparse.ArgumentParser()
    parser.add_argument("schema_folder")
    parser.add_argument("output_dir")
    args = parser.parse_args()
    schema_files = []
    for file in os.listdir(args.schema_folder):
        if file.endswith(".schema"):
            schema_files.append(os.path.join(args.schema_folder, file))
    for schema_file in schema_files:
        print(f"Generating code for {schema_file}")
        generate_cpp_code(schema_file, args.output_dir)
